package ma.zhiqiang.dataline.core.engine;

import cn.hutool.core.date.BetweenFormatter;
import cn.hutool.core.date.DateUtil;
import cn.hutool.core.date.StopWatch;
import cn.hutool.core.util.NumberUtil;
import lombok.extern.slf4j.Slf4j;
import ma.zhiqiang.dataline.core.dto.SinkRecord;
import ma.zhiqiang.dataline.core.dto.SourceRecord;

import java.util.Iterator;
import java.util.concurrent.atomic.AtomicInteger;

@Slf4j
public class Pipeline {
    Source source;
    Transform transform;
    Sink sink;

    public Pipeline(Source source, Transform transform, Sink sink) {
        this.source = source;
        this.transform = transform;
        this.sink = sink;
    }

    Boolean allTaskFinish = false;
    AtomicInteger idx = new AtomicInteger(0);
    AtomicInteger success = new AtomicInteger(0);
    AtomicInteger failed = new AtomicInteger(0);

    public void start() throws Exception{
        StopWatch watch = StopWatch.create("代码执行时间统计");
        try{
            Iterator iterator = source.read();
            watch.start("任务执行开始");
            while (iterator.hasNext()){
                idx.incrementAndGet();
                SourceRecord sourceRecord = source.convert(iterator.next());
                SinkRecord sinkRecord = transform.transfer(sourceRecord);
                Boolean bool = sink.writer(sinkRecord);
                if(bool){
                    success.incrementAndGet();
                }else {
                    failed.incrementAndGet();
                }
            }
        }catch (Exception e){
            log.error("任务执行失败",e);
            throw e;
        }finally {
            if(null!=watch){watch.stop();}
            if(null!=source){source.close();}
        }
        while (!allTaskFinish){
            Thread.sleep(1000);
            if(idx.get()==(success.get()+failed.get())){
                allTaskFinish = Boolean.TRUE;
            }
        }
        String between = DateUtil.formatBetween(watch.getLastTaskTimeMillis(), BetweenFormatter.Level.MILLISECOND);
        Double tps = idx.doubleValue() / watch.getTotalTimeMillis() * 1000;
        log.info("共处理数据[{}]行,成功数量[{}],失败数量[{}],总耗时[{}],TPS[{}/秒]",idx,success,failed,between, NumberUtil.round(tps,2));
    }
}
